package com.wzh.test.avltree;

public class AvlTreeDemo {
    public static void main(String[] args) {
        AvlTree avlTree = new AvlTree();
        avlTree.insert(1);
        avlTree.insert(5);
        avlTree.insert(7);
        avlTree.insert(9);
        avlTree.insert(2);

        avlTree.printTree();
        System.out.println(avlTree.height());
    }
}

class AvlTree{
    private AvlNode root;

    public AvlTree(){
        root = null;
    }

    public void insert(Integer x){
        root = insert(x,root);
    }

    public int height(){
        return height(root);
    }

    private int height(AvlNode tree){
        return tree == null ? -1 : tree.height;
    }

    private AvlNode insert(Integer x,AvlNode tree){
        if(tree == null){
            return new AvlNode(x,null,null);
        }

        int compareResult = x.compareTo(tree.element);

        if(compareResult < 0){
            tree.left = insert(x, tree.left);
            if(height(tree.left) - height(tree.right) == 2){
                if(x.compareTo(tree.left.element) < 0){
                    tree = rotateWithLeftChild(tree);//左-左单旋转
                }else{
                    tree = doubleWithLeftChild(tree);//左-右双旋转
                }
            }
        }else if(compareResult > 0){
            tree.right = insert(x,tree.right);
            if(height(tree.right) - height(tree.left) == 2){
                if(x.compareTo(tree.right.element) > 0){
                    tree = rotateWithRightChild(tree);//右-右单旋转
                }else{
                    tree = doubleWithRightChild(tree);//右-左双旋转
                }
            }
        }

        tree.height = Math.max(height(tree.left),height(tree.right)) + 1;
        return tree;
    }

    private AvlNode rotateWithLeftChild(AvlNode k2){
        AvlNode k1 = k2.left;
        k2.left = k1.right;
        k1.right = k2;

        k2.height = Math.max(height(k2.left),height(k2.right)) + 1;
        k1.height = Math.max(height(k1.left),height(k2.right)) + 1;

        return k1;
    }

    private AvlNode rotateWithRightChild(AvlNode k2){
        AvlNode k1 = k2.right;
        k2.right = k1.left;
        k1.left = k2;

        k1.height = Math.max(height(k1.left),height(k1.right)) + 1;
        k2.height = Math.max(height(k2.left),height(k2.right)) + 1;

        return k1;
    }

    private AvlNode doubleWithLeftChild(AvlNode k3){
        k3.left = rotateWithRightChild(k3.left);
        return rotateWithLeftChild(k3);
    }

    private AvlNode doubleWithRightChild(AvlNode k3){
        k3.right = rotateWithLeftChild(k3.right);
        return rotateWithRightChild(k3);
    }

    public void printTree(){
        printTree(root);
    }

    private void printTree(AvlNode node){
        if(node != null){
            printTree(node.left);
            System.out.println(node.element);
            printTree(node.right);
        }
    }

}
class AvlNode{
    Integer element;
    AvlNode left;
    AvlNode right;
    int height;

    public AvlNode(Integer element){
        this(element,null,null);
    }

    public AvlNode(Integer element,AvlNode left,AvlNode right){
        this.element = element;
        this.left = left;
        this.right = right;
    }


}
